from shapely.geometry import Polygon
import numpy as np
import torch


def iou_polygen(poly_boxes1, poly_boxes2):
    ious = []
    for poly_box1, poly_box2 in zip(poly_boxes1, poly_boxes2):
        if (poly_box1 - poly_box2).sum() == 0:
            ious.append(0.98888)
            continue
        poly1 = Polygon(poly_box1)
        poly2 = Polygon(poly_box2)
        inter_area = poly1.intersection(poly2).area
        union_area = poly1.area + poly2.area - inter_area
        ious.append(inter_area / (union_area))

    ious = np.array(ious)
    return ious.mean()


def iou_polygenMN(poly_boxes1, poly_boxes2):
    """
    poly_boxes1 (N,8)
    poly_boxes2 (M,8)
    return (N,M)
    """
    out = torch.zeros([poly_boxes1.shape[0], poly_boxes2.shape[0]], device=poly_boxes1.device)
    for i, poly_box1 in enumerate(poly_boxes1):
        for j, poly_box2 in enumerate(poly_boxes2):
            if (poly_box1 - poly_box2).sum() == 0:
                out[i, j] = 0.98888
                continue
            poly1 = Polygon(poly_box1.detach().cpu().numpy().reshape([4, 2]))
            poly2 = Polygon(poly_box2.detach().cpu().numpy().reshape([4, 2]))
            try:
                inter_area = poly1.intersection(poly2).area
            except Exception as ex:
                inter_area = 0.0
            union_area = poly1.area + poly2.area - inter_area
            out[i, j] = (inter_area / (union_area + 1e-10))

    return out
